#include "v4l2config.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "mpp_buffer.h"

#define FRAME_COUNT 1000
#define RAW_FILE_NAME "capture.raw"
#define CROPPED_FILE_NAME "cropped.raw"

// 灰度图裁剪并转换为NV12格式
void grey_to_nv12(unsigned char* src_grey, unsigned char* dst_nv12, 
                  int src_width, int src_height, int dst_width, int dst_height) {
    if (!src_grey || !dst_nv12 || src_width <= 0 || src_height <= 0) {
        perror("Invalid buffer or parameters");
        return;
    }

    // 计算裁剪偏移（居中裁剪）
    int offset_x = 0;
    int offset_y = 0;
    
    // 确保偏移量有效
    if (offset_x < 0) offset_x = 0;
    if (offset_y < 0) offset_y = 0;

    // Y分量大小
    const int y_size = dst_width * dst_height;
    // NV12格式中UV分量紧跟Y分量，每个UV对占2字节
    unsigned char* uv_plane = dst_nv12 + y_size;

    // 1. 裁剪并复制Y分量（灰度值直接作为Y）
    for (int y = 0; y < dst_height; ++y) {
        const unsigned char* src_row = &src_grey[(y + offset_y) * src_width + offset_x];
        unsigned char* dst_row = &dst_nv12[y * dst_width];
        memcpy(dst_row, src_row, dst_width);
    }

    // 2. 生成UV分量（灰度图转换为NV12时，UV通常填充为128表示中性值）
    for (int y = 0; y < dst_height / 2; ++y) {
        for (int x = 0; x < dst_width / 2; ++x) {
            int uv_index = y * dst_width + x * 2;
            uv_plane[uv_index] = 128;     // U分量
            uv_plane[uv_index + 1] = 128; // V分量
        }
    }
}

int main() {
    CamSource* cam = nullptr;
    FILE* raw_file = nullptr;
    FILE* cropped_file = nullptr;
    int ret = 0;
    // 提前声明变量，避免goto跳过初始化
    const int dst_width = 640;
    const int dst_height = 480;
    const size_t nv12_size = dst_width * dst_height * 3 / 2;
    unsigned char* nv12_buffer = nullptr;

    printf("Starting camera capture...\n");

    // 初始化摄像头
    cam = camera_source_init("/dev/video0", 4, 640, 512, MPP_FMT_YUV400);
    if (!cam) {
        fprintf(stderr, "Camera initialization failed\n");
        ret = -1;
        goto cleanup;
    }

    // 打开输出文件
    raw_file = fopen(RAW_FILE_NAME, "wb");
    if (!raw_file) {
        fprintf(stderr, "Failed to open raw file: %s\n", RAW_FILE_NAME);
        ret = -1;
        goto cleanup;
    }

    cropped_file = fopen(CROPPED_FILE_NAME, "wb");
    if (!cropped_file) {
        fprintf(stderr, "Failed to open cropped file: %s\n", CROPPED_FILE_NAME);
        ret = -1;
        goto cleanup;
    }

    // 分配NV12输出缓冲区（在所有goto之后）
    nv12_buffer = (unsigned char*)malloc(nv12_size);
    if (!nv12_buffer) {
        fprintf(stderr, "Memory allocation failed for NV12 buffer\n");
        ret = -1;
        goto cleanup;
    }

    // 捕获并处理帧数据
    for (int i = 0; i < FRAME_COUNT; ++i) {
        // 获取帧索引
        int frame_idx = camera_source_get_frame(cam);
        if (frame_idx < 0) {
            fprintf(stderr, "Failed to get frame %d\n", i + 1);
            break;
        }

        // 获取缓冲区
        MppBuffer buf = camera_frame_to_buf(cam, frame_idx);
        if (!buf) {
            fprintf(stderr, "Failed to get buffer for frame %d\n", i + 1);
            camera_source_put_frame(cam, frame_idx);
            continue;
        }

        // 获取原始帧数据
        size_t len = mpp_buffer_get_size(buf);
        void* data = mpp_buffer_get_ptr(buf);
        if (!data || len <= 0) {
            fprintf(stderr, "Invalid buffer data for frame %d\n", i + 1);
            camera_source_put_frame(cam, frame_idx);
            continue;
        }

        // 写入原始数据
        size_t written = fwrite(data, 1, len, raw_file);
        if (written != len) {
            fprintf(stderr, "Warning: Only wrote %zu of %zu bytes for frame %d\n", written, len, i + 1);
        }

        // 转换并写入裁剪后的NV12数据
        grey_to_nv12((unsigned char*)data, nv12_buffer, 768, 512, dst_width, dst_height);
        written = fwrite(nv12_buffer, 1, nv12_size, cropped_file);
        if (written != nv12_size) {
            fprintf(stderr, "Warning: Only wrote %zu of %zu bytes for cropped frame %d\n", written, nv12_size, i + 1);
        }

        printf("Processed frame %d: original %dx%d -> cropped %dx%d\n", 
               i + 1, 768, 512, dst_width, dst_height);

        // 释放缓冲区
        camera_source_put_frame(cam, frame_idx);
    }

    printf("Successfully captured %d frames\n", FRAME_COUNT);

cleanup:
    // 释放资源
    if (nv12_buffer) free(nv12_buffer);
    if (cropped_file) fclose(cropped_file);
    if (raw_file) fclose(raw_file);
    if (cam) camera_source_deinit(cam);

    printf("RAW data saved to %s\n", RAW_FILE_NAME);
    printf("Cropped NV12 data saved to %s\n", CROPPED_FILE_NAME);

    return ret;
}    